#include "main/utils.h"
#include "lib/list.h"
using namespace std;

static bool cmp(ListNode* in1, ListNode* in2) {
  return in1->value > in2->value;
}

ListNode* mergeList(vector<ListNode*> nodes) {
  priority_queue<ListNode*, vector<ListNode*>, decltype(cmp)*> p_queue(cmp);
  for (auto ele : nodes) {
    p_queue.push(ele);
  }
  ListNode* phead = CreateNode(0);
  auto cur_node = phead;
  while (!p_queue.empty()) {
    auto next_node = p_queue.top();
    p_queue.pop();
    cur_node->next = next_node;
    cur_node = cur_node->next;
    next_node = next_node->next;
    if (next_node) {
      p_queue.push(next_node);
    }
  }
  return phead->next;
}

ListNode* merge(ListNode* head1, ListNode* head2) {
  auto phead = CreateNode(0);
  auto node = phead;
  while (head1 && head2) {
    if (head1->value < head2->value) {
      node->next = head1;
      head1 = head1->next;
    } else {
      node->next = head2;
      head2 = head2->next;
    }
    node = node->next;
  }
  node->next = head1 ? head1 : head2;
  return phead->next;
}

ListNode* mergeList(vector<ListNode*> nodes, int begin, int end) {
  if (begin + 1 >= end)
    return nodes[begin];

  int mid = (begin + end) / 2;
  auto head1 = mergeList(nodes, begin, mid);
  auto head2 = mergeList(nodes, mid, end);
  return merge(head1, head2);
}

int main() {
  vector<int> nums1 = {1, 4, 7};
  vector<int> nums2 = {2, 5, 8};
  vector<int> nums3 = {3, 6, 9};
  auto list1 = CreateList(nums1);
  auto list2 = CreateList(nums2);
  auto list3 = CreateList(nums3);
  vector<ListNode*> nodes = {list1, list2, list3};
  auto ret = mergeList(nodes);
  PrintList(ret);

  auto list4 = CreateList(nums1);
  auto list5 = CreateList(nums2);
  auto list6 = CreateList(nums3);
  nodes = {list4, list5, list6};
  ret = mergeList(nodes, 0, nodes.size());
  PrintList(ret);

  return 0;
}
